
#### block boot ####
tjbal_block_boot = function(data, unit, time, 
                            outcome, treated, nboot){
  
  set.seed(773212)
  
  lapply(c('dplyr', 'purrr', 'tjbal'), require, character.only = T)
  
  quiet <- function(x) { 
    sink(tempfile()) 
    on.exit(sink()) 
    invisible(force(x)) 
  } 
  
  data = data %>% 
    select(unit, time, treated, outcome) 
  colnames(data)[1:4] = c('unit', 'time', 'treated', 'outcome')
  
  baseline_year = data %>% 
    group_by(time) %>% 
    summarise(na_ever = mean(outcome, na.rm = T), 
              is_na_year = as.numeric(is.na(na_ever) == FALSE, 1,0)) %>% 
    filter(is_na_year == 1) %>% 
    summarise(min_year = min(time))
  
  baseline_year = as.numeric(baseline_year)
  
  force_balance = function(data, outcome, baseline_year){
    
    data <- data %>% 
      group_by(unit) %>% 
      filter(time >= baseline_year) %>% 
      mutate(na_ever = mean(outcome)) %>% 
      filter(is.na(na_ever) == F) %>% 
      select(-c(na_ever)) %>% 
      as.data.frame()
    
    return(data)
  }
  
  data = force_balance(data = data, outcome = outcome, 
                       baseline_year = baseline_year)
  data = data %>% 
    group_by(time) %>% 
    mutate(NA_for_year = as.numeric(is.na(mean(outcome)) == F, 1,0)) %>% 
    filter(NA_for_year == 1) %>% 
    as.data.frame()
  
  
  tjbal_boot = function(data, unit, time, 
                        outcome, treated, baseline_year){
    
    get_block_index <- function(n_block, # Number of block
                                size # Number of observation in each block
    ) {
      
      block1 <- 1:size
      
      step <- 0:(n_block-1)*size
      step_boot <- sample(step, n_block, replace = TRUE)
      
      step_expand <- rep(step_boot, each = size)
      block_expand <- rep(block1, n_block)
      
      block_expand + step_expand
      
    }
    
    n_block = length(unique(data$unit)) 
    size = length(unique(data$time))
    
    block_index = get_block_index(n_block = n_block, size)
    block_data = data[block_index, ]$unit %>% unique()
    block_data = as.data.frame(cbind(block_data,1))
    colnames(block_data)[1:2] <- c('unit', 'sample')
    df2 <- merge(data, block_data, by = 'unit', all.x = T)
    
    df2 <- df2 %>% 
      filter(sample == 1)
    
    m = quiet(tjbal(data = df2, Y = 'outcome', D = "treated", 
                    index = c("unit", "time"), 
                    demean = F, estimator = "mean", 
                    vce = "none"))
    att_avg = m$att.avg
    return(att_avg )
  }
  
  map_progress <- function(.x, .f, ..., .id = NULL) {
    .f <- purrr::as_mapper(.f, ...)
    pb <- progress::progress_bar$new(total = length(.x), format = " [:bar] :current/:total (:percent) eta: :eta", force = TRUE)
    
    f <- function(...) {
      pb$tick()
      .f(...)
    }
    purrr::map(.x, f, ..., .id = .id)
  }
  
  result_list = map_progress(seq_len(nboot), ~
                               tjbal_boot(data, 'unit', 'time', 'outcome', 'treat', baseline_year)
  )
  
  boot_se = as.vector(unlist(result_list))
  boot_se = sd(boot_se)
  
  m = tjbal(data = data, Y = 'outcome', D = "treated", 
            index = c("unit", "time"), 
            demean = F, estimator = "mean", 
            vce = "none")
  att  = m$att.avg
  
  t_stat = att/boot_se
  
  n_unit = length(unique(data$unit))
  return(c(att, boot_se, t_stat, n_unit))
}

#### het fx #### 
het_fx = function(data, moderator, outcome, year_baseline, max_year){
  
  data = data 
  data$outcome = scale(outcome) 
  data$moderator = moderator  
  max_year = max_year
  
  z_stat_calc <- function(coef1, se1, coef2, se2){ 
    (coef1-coef2)/sqrt(se1^2 + se2^2)
  }
  
  data_one = data %>% 
    filter(is.na(cohort) == F) %>% 
    filter(moderator == 1) %>% 
    filter(year <= 2014) %>%
    filter(year >= year_baseline) %>% 
    mutate(
      treat = ifelse(treated_ever == 1, as.numeric(year >= cohort, 1,0), 0)
    ) %>% 
    group_by(DISTID, year) %>% 
    summarise(outcome = mean(outcome, na.rm = T), 
              treat = mean(treat, na.rm = T)
    ) %>% 
    group_by(DISTID) %>% 
    mutate(NA_data = mean(outcome), 
           n_year = n()
    ) %>% 
    filter(is.na(NA_data ) == F) %>% 
    filter(n_year == max_year) %>% 
    as.data.frame()
  
  data_two = data %>% 
    filter(is.na(cohort) == F) %>% 
    filter(moderator == 0) %>% 
    filter(year <= 2014) %>%
    filter(year >= year_baseline) %>% 
    mutate(
      treat = ifelse(treated_ever == 1, as.numeric(year >= cohort, 1,0), 0)
    ) %>% 
    group_by(DISTID, year) %>% 
    summarise(outcome = mean(outcome, na.rm = T), 
              treat = mean(treat, na.rm = T)
    ) %>% 
    group_by(DISTID) %>% 
    mutate(NA_data = mean(outcome), 
           n_year = n()
    ) %>% 
    filter(is.na(NA_data ) == F) %>% 
    filter(n_year == max_year) %>% 
    as.data.frame()
  
  
  df_one =  
    print(tjbal(data = data_one, 
                Y = 'outcome', 
                D = "treat", 
                index = c("DISTID", "year"), 
                demean = F, 
                estimator = "mean", 
                vce = "jackknife"))[c(1,2,6)]
  
  df_two =  
    print(tjbal(data = data_two, 
                Y = 'outcome', 
                D = "treat", 
                index = c("DISTID", "year"), 
                demean = F, 
                estimator = "mean", 
                vce = "jackknife"))[c(1,2,6)]
  
  z_stat_out = z_stat_calc(coef1 = df_one[1], se1 = df_one[2], 
                           coef2 = df_two[1], se2 =  df_two[2])
  diff_out = df_one[1] - df_two[1]
  
  p_value = 2*pnorm(abs(z_stat_out), lower.tail = F) 
  
  return(cbind(diff_out, z_stat_out, p_value))
  
}


#### make att plot #### 
make_tj_att <- function(out0, xlab2, ylab2, title2){
  
  
  pgks <- c('dplyr', "ggplot2", 'ggthemes', "cowplot", 'tjbal')
  lapply(pgks, require, character.only = T)
  
  out0 <- out0 
  plotdf <- plot(out0, type = "counterfactual", count = FALSE)$data
  plotdf$time <- plotdf$time-1
  
  
  xlab2 = xlab2
  ylab2 = ylab2
  title2 = title2 
  
  plotdf <- plot(out0)$data
  plotdf$time <- plotdf$time-1
  
  m <- print(out0)
  
  
  
  tj_ATT <- ggplot(plotdf, aes(time, ATT)) + 
    geom_line(size = 1.2) +   
    geom_ribbon(aes(ymin=CI.lower, ymax=CI.upper), alpha=0.2) + 
    scale_x_continuous(n.breaks = 6) +
    geom_vline(xintercept = -1, col = 'red', lty = 2) + 
    ggtitle(title2) + 
    xlab(xlab2) + 
    geom_hline(yintercept = 0, col = 'black') + 
    geom_label(data = plotdf, aes(x = max(time)-1, y = max(CI.upper)+.02, 
                                  label = paste(paste("ATT:",  round(m[1], 3), sep = " "), 
                                                paste("SE:", round(m[2],3)))), family = 'Times') 
  return(tj_ATT)
  
}

#### placebo se #### 

placebo_se = function(data, unit, time, treated, cohort, 
                      outcome, nsims, estimator, beep, X, att_only){
  
  if("dplyr" %in% rownames(installed.packages()) == FALSE) {install.packages("dplyr")}
  if("tjbal" %in% rownames(installed.packages()) == FALSE) {install.packages("tjbal")}
  if("purrr" %in% rownames(installed.packages()) == FALSE) {install.packages("purrr")}
  if("beepr" %in% rownames(installed.packages()) == FALSE) {install.packages("beepr")}
  
  pgks <- c('dplyr','tjbal', 'purrr', 'randomizr')
  lapply(pgks, require, character.only = T)
  
  set.seed(303073)
  
  print("Inferring...")
  
  quiet <- function(x) { 
    sink(tempfile()) 
    on.exit(sink()) 
    invisible(force(x)) 
  } 
  
  #### prep data #### 
  data = data %>% 
    select(unit, time, outcome, treated, cohort, X)
  
  colnames(data)[1:5] = c('unit', 'time', 'outcome', 'treated', 'cohort')
  
  baseline_year = data %>% 
    group_by(time) %>% 
    summarise(na_ever = mean(outcome, na.rm = T), 
              is_na_year = as.numeric(is.na(na_ever) == FALSE, 1,0)) %>% 
    filter(is_na_year == 1) %>% 
    summarise(min_year = min(time))
  
  baseline_year = as.numeric(baseline_year)
  
  if (att_only == TRUE){
    placebo_treatment <- function(data, outcome, 
                                  baseline_year, 
                                  X,
                                  estimator){
      
      force_balance = function(data, outcome, baseline_year){
        
        
        data <- data %>% 
          group_by(unit) %>% 
          filter(time >= baseline_year) %>% 
          mutate(na_ever = mean(outcome)) %>% 
          filter(is.na(na_ever) == F) %>% 
          select(-c(na_ever)) %>% 
          as.data.frame()
        
        return(data)
      }
      
      data = force_balance(data = data, outcome = outcome, 
                           baseline_year = baseline_year
      )
      list_of_cohorts = unique(data$cohort[data$treated==1])
      
      size_list = list()
      
      for (j in 1:length(list_of_cohorts)) {
        
        size_list[j] = data %>% filter(time == min(time)) %>%  
          filter(cohort == list_of_cohorts[j]) %>%  
          summarise(n = n())
        
      }
      
      name_cohorts = unique(data$cohort)
      
      df_control_cs = data %>% 
        filter(cohort == 0) %>%
        filter(time == min(time)) %>% 
        select(unit)
      
      n_control = length(unique(df_control_cs$unit))-sum(unlist(size_list))
      
      df_control_cs$g = complete_ra(N = length(unique(df_control_cs$unit)), 
                                    m_each = unlist(c(n_control, size_list)  ), 
                                    conditions = name_cohorts
      )
      data = data %>%  
        filter(cohort == 0)
      data <- merge(data, df_control_cs, by = 'unit', all.x = T)
      data$g = as.character(data$g)
      
      data$treat <- ifelse(as.numeric(data$g) > 0, 
                           as.numeric(data$time>=as.numeric(data$g), 1,0), 0)
      
      m = quiet(tjbal(data, Y = 'outcome', 
                      D = "treat", 
                      index = c("unit", "time"), 
                      X = X,
                      demean = F, 
                      estimator = estimator, 
                      vce = 'none'))
      att_avg = m$att.avg
      return(att_avg)
      
    }
    
    map_progress <- function(.x, .f, ..., .id = NULL) {
      .f <- purrr::as_mapper(.f, ...)
      pb <- progress::progress_bar$new(total = length(.x), format = " [:bar] :current/:total (:percent) eta: :eta", force = TRUE)
      
      f <- function(...) {
        pb$tick()
        .f(...)
      }
      purrr::map(.x, f, ..., .id = .id)
    }
    
    result_list <- unlist(map_progress(seq_len(nsims), 
                                       ~placebo_treatment(data, outcome, baseline_year, 
                                                          X, estimator)
    ))
    
    if(beep == TRUE){
      beepr::beep(sound = 2)
      return(result_list)
    }
    else{return(result_list)}
    
  }
  
  else{
    placebo_treatment <- function(data, outcome, 
                                  baseline_year, 
                                  X,
                                  estimator){
      
      force_balance = function(data, outcome, baseline_year){
        
        data <- data %>% 
          group_by(unit) %>% 
          filter(time >= baseline_year) %>% 
          mutate(na_ever = mean(outcome)) %>% 
          filter(is.na(na_ever) == F) %>% 
          select(-c(na_ever)) %>% 
          as.data.frame()
        
        return(data)
      }
      
      data = force_balance(data = data, outcome = outcome, 
                           baseline_year = baseline_year
      )
      list_of_cohorts = unique(data$cohort[data$treated==1])
      
      size_list = list()
      
      for (j in 1:length(list_of_cohorts)) {
        
        size_list[j] = data %>% filter(time == min(time)) %>%  
          filter(cohort == list_of_cohorts[j]) %>%  
          summarise(n = n())
        
      }
      
      name_cohorts = unique(data$cohort)
      
      df_control_cs = data %>% 
        filter(cohort == 0) %>%
        filter(time == min(time)) %>% 
        select(unit)
      
      n_control = length(unique(df_control_cs$unit))-sum(unlist(size_list))
      
      df_control_cs$g = complete_ra(N = length(unique(df_control_cs$unit)), 
                                    m_each = unlist(c(n_control, size_list)  ), 
                                    conditions = name_cohorts
      )
      data = data %>%  
        filter(cohort == 0)
      data <- merge(data, df_control_cs, by = 'unit', all.x = T)
      data$g = as.character(data$g)
      
      data$treat <- ifelse(as.numeric(data$g) > 0, 
                           as.numeric(data$time>=as.numeric(data$g), 1,0), 0)
      
      m = quiet(tjbal(data, Y = 'outcome', 
                      D = "treat", 
                      index = c("unit", "time"), 
                      X = X,
                      demean = F, 
                      estimator = estimator, 
                      vce = 'none'))
      att_ts = m$att
      return(att_ts)  
    }
    
    map_progress <- function(.x, .f, ..., .id = NULL) {
      .f <- purrr::as_mapper(.f, ...)
      pb <- progress::progress_bar$new(total = length(.x), format = " [:bar] :current/:total (:percent) eta: :eta", force = TRUE)
      
      f <- function(...) {
        pb$tick()
        .f(...)
      }
      purrr::map(.x, f, ..., .id = .id)
    }
    
    result_list <- unlist(map_progress(seq_len(nsims), 
                                       ~placebo_treatment(data, outcome, baseline_year, 
                                                          X, estimator)
    ))
    
    sim_id = 1:nsims
    event_time_min = baseline_year-(max(as.numeric(data$cohort)[data$treat==1])-1)
    event_time_max = max(data$time)-(min(as.numeric(data$cohort)[data$treat==1])-1)
    event_times = seq(event_time_min, event_time_max , 1)
    
    result_data_frame = merge(sim_id, event_times)
    result_data_frame = cbind.data.frame(result_data_frame, result_list)
    
    force_balance = function(data, outcome, baseline_year){
      
      data <- data %>% 
        group_by(unit) %>% 
        filter(time >= baseline_year) %>% 
        mutate(na_ever = mean(outcome)) %>% 
        filter(is.na(na_ever) == F) %>% 
        select(-c(na_ever)) %>% 
        as.data.frame()
      
      return(data)
    }
    
    data = force_balance(data = data, outcome = outcome, 
                         baseline_year = baseline_year
    )
    
    m_true = tjbal(data, Y = 'outcome', 
                   D = "treated", 
                   index = c("unit", "time"), 
                   X = X,
                   demean = F, 
                   estimator = estimator, 
                   vce = 'none')
    
    att_ts = cbind.data.frame(0, m_true$att)
    att_ts$y = row.names(att_ts)
    colnames(att_ts)[1:3] = c('x', 'result_list',  'y')
    
    result_data_frame = rbind.data.frame(att_ts, result_data_frame)
    result_data_frame$y = as.numeric(result_data_frame$y)
    
    if(beep == TRUE){
      beepr::beep(sound = 2)
      return(result_data_frame)
    }
    else{return(result_data_frame)}
    
  }
}

make_placebo_plot = function(placebo_estimate){
  
  plot = ggplot(placebo_estimate, aes(y, result_list)) + 
    geom_point(data = subset(placebo_estimate, x > 0), 
               alpha = .4, size = 2, position = "jitter", color = 'blue') + 
    geom_point(data = subset(placebo_estimate, x == 0), 
               alpha = 1, size = 4, color = 'red') + 
    geom_line(data = subset(placebo_estimate, x == 0), 
              alpha = .4, size = 2,  color = 'red') + 
    xlab('Time until Courts') + 
    ylab('ATT') + 
    scale_color_manual() + 
    geom_vline(xintercept = 0, col = 'red', lty = 2) + 
    theme_bw()  
  return(plot)
  
}


#### make tj ct #### 

make_tj_ct = function(out0, xlab1, ylab1, title1){
  
  pgks <- c('dplyr', "ggplot2", 'ggthemes', "cowplot", 'tjbal')
  lapply(pgks, require, character.only = T)
  
  
  
  ybar_df <- as.data.frame(out0$Y.bar[,1])
  ybar_df$time <- row.names(ybar_df)
  ybar_df$time <- as.numeric(ybar_df$time)-1 
  ybar_df$type <- 'treat'
  colnames(ybar_df)[1] <- 'outcome'
  
  ct_df <- as.data.frame(out0$Y.bar[,2])
  ct_df$time <- row.names(ct_df)
  ct_df$time <- as.numeric(ct_df$time)-1 
  ct_df$type <- 'ct'
  colnames(ct_df)[1] <- 'outcome'
  
  
  plotdf <- rbind.data.frame(ybar_df, ct_df)
  
  xlab1 = xlab1
  ylab1 = ylab1
  title1 = title1 
  
  
  tj_trend <- ggplot(plotdf, aes(time, outcome, lty = type, color = type)) +
    geom_line(size = 1) + 
    geom_point() + 
    scale_x_continuous(n.breaks = 6) +
    ggtitle(title1) + 
    labs(lty='Courts', color = 'Courts') +
    xlab(xlab1) + 
    scale_color_manual(values = c('red', 'black'), labels = c('Counterfactual', 'Courts')) + 
    scale_linetype_manual(values = c('dashed', 'solid'), labels = c('Counterfactual', 'Courts'))  +
    geom_vline(xintercept = -1, col = 'red', lty = 2) + 
    ylab(ylab1) + 
    theme(legend.position = 'bottom') + 
    theme(text = element_text(size=20)) 
  
  return(tj_trend)
  
  
}

#### recode levels #### 

recodr <- function(x){ 
  
  x <- as.character(x)
  x[x=="Very well"] <- 5
  x[x=="A little well"] <- 4
  x[x=="Neither poorly, nor well"] <- 3
  x[x=="A little poorly"] <- 2
  x[x=="Very poorly"] <- 1
  x <- ifelse(x>0, as.numeric(x), NA)
}

#### tjbal unbalanced #### 

tjbal_unbalance = function(data, outcome, baseline_year, 
                           X,
                           estimator_type, vce){
  
  force_balance = function(data, outcome, baseline_year){
    
    data = data 
    data$outcome = outcome 
    
    data <- data %>% 
      group_by(DISTID) %>% 
      filter(year >= baseline_year) %>% 
      mutate(na_ever = mean(outcome)) %>% 
      filter(is.na(na_ever) == F) %>% 
      as.data.frame()
    
    return(data)
  }
  
  data = force_balance(data, outcome, baseline_year)
  
  m = tjbal(data, Y = 'outcome', 
            D = "treat", 
            index = c("DISTID", "year"), 
            X = X,
            demean = F, estimator = estimator_type, 
            vce = vce)
  
  return(m)
}



#### table #### 

make_tjbal_table <- function(att_tjbal, round_to_nearest, n_models_per_out, 
                             out_names, data, id, time, treated_group, outcome_matrix, header, 
                             header_labels, label_treatment, extra_arg, file_save){
  
  packs <- c('xtable',  'tidyr', 'dplyr', 'purrr')
  suppressMessages(lapply(packs, require, character.only = T))
  
  data <- data 
  file_save <- file_save
  treated_group <- treated_group
  
  n_unit <- length(unique(id))
  n_time <- length(unique(time)) 
  
  J <- length(out_names)
  
  att_tjbal <- as.data.frame(att_tjbal)
  colnames(att_tjbal)[1:3] <- c('att', 'se', 'pval')
  att_tjbal$att  <- round(att_tjbal$att, round_to_nearest)
  att_tjbal$se  <- round(att_tjbal$se, round_to_nearest)
  
  outcome_data <- outcome_matrix
  round_to_nearest_fun <- function(x){
    round(x, round_to_nearest)
  }
  
  sd_dv <- map(map(outcome_data, sd), round_to_nearest_fun)
  
  mean_dv_data <- cbind.data.frame(outcome_matrix, treated_group)
  mean_dv_data <- mean_dv_data %>% 
    filter(treated_group == 0)
  mean_dv_control <- map(map(mean_dv_data, mean), round_to_nearest_fun)
  
  
  att_tjbal$outcome_name <- seq(1, J, 1)
  colnames(att_tjbal)[1:3] <- c('att', 'se', 'pval')
  att_tjbal <- att_tjbal %>% 
    mutate(  att = if_else(pval < .1 & pval >= .05, paste0("$", att, "^{\\dagger}$"), paste0(att)), 
             att = if_else(pval < .05 & pval >= .01, paste0("$", att, "^{*}$"), paste0(att)), 
             att = if_else(pval < .01 & pval >= .001, paste0("$", att, "^{**}$"), paste0(att)), 
             att = if_else(pval < .001, paste0("$", att, "^{***}$"), paste0(att)), 
             se = paste0("(", se,")")
    ) %>% 
    mutate_if(is.numeric, as.character)
  
  att_tjbal <- att_tjbal %>% 
    pivot_longer(cols = c('att', 'se', 'pval'))
  
  att_tjbal <- att_tjbal %>% 
    pivot_wider(names_from = outcome_name, 
                values_from = value)
  att_tjbal <- att_tjbal[1:2,]
  att_tjbal[1,1] <- label_treatment
  att_tjbal[2,1] <- ""
  
  colnames(att_tjbal)[1:ncol(att_tjbal)] <- ""
  colnames(att_tjbal)[2:ncol(att_tjbal)]  <- paste0("(", seq(1, J, 1), ")")
  
  att_tjbal <- rbind.data.frame(att_tjbal,
                                dataset = extra_arg, 
                                Districts  = c("N. Districts", rep(n_unit, J)),
                                Years  = c("N. Years", rep(n_time, J)),
                                Stdev  = c("Standard Deviation DV", unlist(sd_dv)
                                ),
                                MeanDV  = c("Mean DV (Control)", unlist(mean_dv_control)
                                )
  )
  
  if(header == T){
    
    first_paste_arg <- paste0("& \\multicolumn{", n_models_per_out, "}{c}{")
    
    addtorow <- list()
    addtorow$pos <- list(0)
    addtorow$command <- paste0(
      
      paste0(
        paste0(paste0(first_paste_arg, unique(header_labels), '}', collapse='')),
        "\\\\"),
      
      paste(paste('Outcome',
                  gsub(",", "&", toString(out_names)),sep = " &"), "\\\\")
      
      
    )
    
    
    return(print(xtable(att_tjbal), add.to.row=addtorow,type="latex",include.rownames=F,
                 booktabs = T, sanitize.text.function=identity, file = file_save))
  }
  else{
    
    addtorow <- list()
    addtorow$pos <- list(0)
    addtorow$command <- paste(paste('Outcome',
                                    gsub(",", "&", toString(out_names)),sep = " &"), "\\\\")
    return(print(xtable(att_tjbal), add.to.row=addtorow,type="latex",include.rownames=F,
                 booktabs = T, sanitize.text.function=identity, file = file_save))
  }
  
}


